package com.starhub.application.rag.service;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import com.starhub.application.rag.biz.MilvusClientFactory;
import com.starhub.application.rag.config.MilvusEmbeddingConfig;
import com.starhub.application.rag.store.LangchainStoreBuilder;
import com.starhub.common.bean.model.MilvusConfig;
import com.starhub.application.rag.constant.RagConstant;

import io.milvus.client.MilvusServiceClient;
import io.milvus.param.R;
import io.milvus.param.collection.FlushParam;
import io.milvus.grpc.FlushResponse;
import lombok.extern.slf4j.Slf4j;

/**
 * Milvus 刷新服务
 * 用于手动管理 Milvus 集合的刷新操作，避免速率限制
 * 
 * @version: V1.0
 * @author: liuhf
 */
@Slf4j
@Service
public class MilvusFlushService {

    @Autowired
    private MilvusEmbeddingConfig milvusEmbeddingConfig;

    @Autowired
    private MilvusConfigService milvusConfigService;

    @Autowired
    private MilvusClientFactory milvusClientFactory;

    @Autowired
    private LangchainStoreBuilder langchainStoreBuilder;

    // 记录每个集合的最后刷新时间
    private final Map<String, Long> lastFlushTime = new ConcurrentHashMap<>();
    
    // 定时刷新任务执行器
    private final ScheduledExecutorService flushScheduler = Executors.newSingleThreadScheduledExecutor();

    /**
     * 手动刷新指定集合
     * 
     * @param collectionName 集合名称
     * @return 是否成功
     */
    public boolean flushCollection(String collectionName) {
        try {
            MilvusConfig config = milvusConfigService.getConfigByMark(RagConstant.MILVUS_MARK);
            MilvusServiceClient client = milvusClientFactory.getClient(RagConstant.MILVUS_MARK);

            FlushParam flushParam = FlushParam.newBuilder()
                    .withCollectionNames(java.util.Collections.singletonList(collectionName))
                    .build();

            R<FlushResponse> response = client.flush(flushParam);
            
            if (response.getStatus() == R.Status.Success.getCode()) {
                lastFlushTime.put(collectionName, System.currentTimeMillis());
                log.info("Successfully flushed collection: {}", collectionName);
                return true;
            } else {
                log.error("Failed to flush collection: {}, error: {}", collectionName, response.getMessage());
                return false;
            }
        } catch (Exception e) {
            log.error("Error flushing collection: {}", collectionName, e);
            return false;
        }
    }

    /**
     * 检查是否需要刷新集合
     * 
     * @param collectionName 集合名称
     * @return 是否需要刷新
     */
    public boolean needsFlush(String collectionName) {
        Long lastFlush = lastFlushTime.get(collectionName);
        if (lastFlush == null) {
            return true;
        }
        
        long currentTime = System.currentTimeMillis();
        return (currentTime - lastFlush) >= milvusEmbeddingConfig.getFlushIntervalMs();
    }

    /**
     * 启动定时刷新任务
     */
    public void startScheduledFlush() {
        flushScheduler.scheduleAtFixedRate(() -> {
            try {
                // 这里可以添加逻辑来获取所有需要刷新的集合
                // 目前先记录日志
                log.debug("Scheduled flush task running...");
            } catch (Exception e) {
                log.error("Error in scheduled flush task", e);
            }
        }, milvusEmbeddingConfig.getFlushIntervalMs(), milvusEmbeddingConfig.getFlushIntervalMs(), TimeUnit.MILLISECONDS);
        
        log.info("Started scheduled flush task with interval: {} ms", milvusEmbeddingConfig.getFlushIntervalMs());
    }

    /**
     * 停止定时刷新任务
     */
    public void stopScheduledFlush() {
        if (!flushScheduler.isShutdown()) {
            flushScheduler.shutdown();
            try {
                if (!flushScheduler.awaitTermination(60, TimeUnit.SECONDS)) {
                    flushScheduler.shutdownNow();
                }
            } catch (InterruptedException e) {
                flushScheduler.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }
        log.info("Stopped scheduled flush task");
    }

    /**
     * 获取集合的最后刷新时间
     * 
     * @param collectionName 集合名称
     * @return 最后刷新时间戳
     */
    public Long getLastFlushTime(String collectionName) {
        return lastFlushTime.get(collectionName);
    }
} 